from langchain_openai import ChatOpenAI
from globle import envs
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate

schema = {'金融':  ['日期', '股票名称', '开盘价', '收盘价', '成交量']}
sentences = ["2023-02-15，寓意吉祥的节日，股票佰笃[BD]美股开盘价10美元，虽然经历了波动，但最终以13美元收盘，成交量微幅增加至460,000，投资者情绪较为平稳",
"2023-04-05，市场迎来轻松氛围，股票盘古(0021)开盘价23元，尽管经历了波动，但最终以26美元收盘，成交量缩小至310,000，投资者保持观望态度。"]

'''
需求：根据这个schema抽取两句话中的金融信息，分开抽取，返回json
'''

system_prompt = '''你是一个信息抽取助手，信息schema如下：
    '金融':  ['日期', '股票名称', '开盘价', '收盘价', '成交量']
    根据schema抽取用户提供的信息，返回json
'''
# 注意：{}是变量占位符，如果不需要变量，不要随意写在prompt里
prompt = ChatPromptTemplate.from_messages(
    [('system',system_prompt),
     ('user','{text}')]
)

model = ChatOpenAI(model="deepseek-chat",  # 或具体模型名
                   base_url="https://api.deepseek.com",  # DeepSeek API 地址,
                    temperature=0.5)

chain = prompt | model| JsonOutputParser()
for sent in sentences:
    r = chain.invoke({"text":sent})
    print(r)